



# 本地库
from models import GGNN, residual_graph_attention, GNN_FiLM, Edge_Conv, MTFF_Co_Attention, Tensor_GCN
from models import Transformer_GCN, Relational_GCN, Deep_GCN






concat_singal = "_- ,/\\"






def name_to_model(name: str, args, **kwargs):
    """将字符串转为模型并返回。

    Args:
        name (str): 模型名称。推荐小写。
        args (_type_): 可能会用到的参数
        **kwargs: 可能需要传入的其他字典型参数。

    Raises:
        ValueError: 类名不存在
    """
    name = name.lower()
    name = name.replace(concat_singal, "")
    if name in ["ggnn", "graphgatedneuralnetwork"]:
        return GGNN(num_edge_types=args.num_edge_types,
                    in_features=args.graph_node_max_num_chars,
                    out_features=args.out_features,
                    embedding_out_features=args.h_features,
                    embedding_num_classes=70,
                    dropout=args.dropout_rate,
                    device=args.device)
    elif name in ["resgagn"]:
        return residual_graph_attention(num_edge_types=args.num_edge_types,
                       in_features=args.graph_node_max_num_chars,
                       out_features=args.out_features,
                       embedding_out_features=args.h_features,
                       embedding_num_classes=70,
                       dropout=args.dropout_rate,
                       max_node_per_graph=args.max_node_per_graph,
                       device=args.device)
    elif name in ["gnn_film", "gnnfilm"]:
        return GNN_FiLM(num_edge_types=args.num_edge_types,
                        in_features=args.graph_node_max_num_chars,
                        out_features=args.out_features,
                        embedding_out_features=args.h_features,
                        embedding_num_classes=70,
                        dropout=args.dropout_rate,
                        device=args.device)
    elif name in ["edge_conv", "edgeconv"]:
        return Edge_Conv(num_edge_types=args.num_edge_types,
                         in_features=args.graph_node_max_num_chars,
                         out_features=args.out_features,
                         embedding_out_features=args.h_features,
                         embedding_num_classes=70,
                         dropout=args.dropout_rate,
                         device=args.device)
    elif name in ["mtff_co_attention", "mtff_module", "mtff_model", "mtff"]:
        return MTFF_Co_Attention(
            feature_models=kwargs["feature_models"],
            in_features=args.graph_node_max_num_chars,
            out_features=args.out_features,
            embedding_out_features=args.h_features,
            embedding_num_classes=70,
            max_node_per_graph=args.max_node_per_graph,
            device=args.device
        )
    elif name in ["tensorgcn", "tensor_gcn"]:
        return Tensor_GCN(
            num_edge_types=args.num_edge_types,
            in_features=args.graph_node_max_num_chars,
            out_features=args.out_features,
            embedding_out_features=args.h_features,
            embedding_num_classes=70,
            dropout=args.dropout_rate,
            max_node_per_graph=args.max_node_per_graph,
            device=args.device
        )
    elif name in ["transformer_gcn", "transformergcn", "tgcn"]:
        return Transformer_GCN(num_edge_types=args.num_edge_types,
                    in_features=args.graph_node_max_num_chars,
                    out_features=args.out_features,
                    embedding_out_features=args.h_features,
                    embedding_num_classes=70,
                    dropout=args.dropout_rate,
                    device=args.device)
    elif name in ["realtional_gcn", "realtionalgcn", "rgcn"]:
        return Relational_GCN(num_edge_types=args.num_edge_types,
                    in_features=args.graph_node_max_num_chars,
                    out_features=args.out_features,
                    embedding_out_features=args.h_features,
                    embedding_num_classes=70,
                    dropout=args.dropout_rate,
                    device=args.device)
    elif name in ["deep_gcn", "deepgcn", "dgcn"]:
        return Deep_GCN(num_edge_types=args.num_edge_types,
                    in_features=args.graph_node_max_num_chars,
                    out_features=args.out_features,
                    embedding_out_features=args.h_features,
                    embedding_num_classes=70,
                    dropout=args.dropout_rate,
                    device=args.device)
    else:
        raise ValueError("Unkown model name '%s'" % name)


